<script setup lang="ts">
import { Ref, computed, onMounted, ref, watch } from "vue";
import ReviewScene from "./components/ReviewScene.vue";
import { invoke } from "@tauri-apps/api";
import { Quaternion } from "@babylonjs/core";
import { loadFrames } from "./utils";


const frameIndex = ref(0)
const totalFrame = ref(0)
const onlyShowUnmark = ref(false);
const playing = ref(false)

const markings: Ref<Map<number, number>> = ref(new Map());

const stage: Ref<HTMLCanvasElement | null> = ref(null)

const videoFrames: Ref<Array<ImageBitmap> | null> = ref(null)

const videoIndexes: Ref<Array<number>> = ref([])
const currVideoId: Ref<number | null> = ref(null)

const playInterval: Ref<number | null> = ref(null)

const targetPos: Ref<Quaternion | null> = ref(null)
const predRot: Ref<Quaternion | null> = ref(null)

const stopPlay = () => {
  const id = playInterval.value;
  if (id !== null) {
    clearInterval(id)
    playInterval.value = null;
  }
  playing.value = false
}

interface Rot {
  raw: Array<number>,
  pred: Array<number> | null,
}

const renderStage = (frameIdx: number) => {
  if (videoFrames.value !== null && frameIdx <= videoFrames.value.length) {
    if (stage.value !== null) {
      const ctx = stage.value.getContext('2d')
      ctx?.drawImage(videoFrames.value[frameIdx - 1], 0, 0)
    }
  }
}

const updateRealPos = async (curr: number) => {
  if (currVideoId.value !== null) {
    try {
      // imageUrl.value = `frame://localhost/${videoId}/${curr}`
      const rot: Rot = await invoke('frame_raw', {videoId: currVideoId.value, frameId: curr})
      const quat = rot.raw;
      const q = new Quaternion(
        -quat[2],  // x := -y
        -quat[3],  // y := -z
        -quat[1],  // z := -x
        -quat[0]
      )
      probeRealQuat.value = q

      if (rot.pred !== null) {
        const pred = new Quaternion(
          -rot.pred[2],  // x := -y
          -rot.pred[3],  // y := -z
          -rot.pred[1],  // z := -x
          -rot.pred[0]
        )
        predRot.value = q.multiply(pred)
      } else {
        predRot.value = null
      }
    } catch(e) {
      console.error(e)
    }
  }
}

const updateTargetPos = async (videoId: number, frameCount: number) => {
  const rot: Rot = await invoke('frame_raw', {videoId: videoId, frameId: frameCount})
  const quat = rot.raw;
  const q = new Quaternion(
    -quat[2],  // x := -y
    -quat[3],  // y := -z
    -quat[1],  // z := -x
    -quat[0]
  )


  targetPos.value = q;
}

const updateVideoId = async (content: string) => {
  let id = parseInt(content, 10);
  if (videoIndexes.value.find((i) => i === id) === undefined) {
    return;
  }
  stopPlay()
  frameIndex.value = 1
  currVideoId.value = id;
  const frameCount: number = await invoke('video_frame_count', { id })

  console.log("begin load frame")
  const frames = await loadFrames(id, frameCount);
  videoFrames.value = frames;

  totalFrame.value = frameCount;
  predRot.value = null;
  renderStage(frameIndex.value)
  updateRealPos(frameIndex.value)
  updateTargetPos(currVideoId.value, frameCount)
}

const nextFrame = () => {
  if (frameIndex.value < totalFrame.value) {
    frameIndex.value += 1
  }
}

const prevFrame = () => {
  console.log('prev frame')
  if (frameIndex.value > 1) {
    frameIndex.value -= 1
  }
}

const toKeyFrame = () => {
  console.log('to keyframe')
  if (keyFrameMarked.value && currVideoId.value !== null) {
    const frame = markings.value.get(currVideoId.value)
    if (frame !== undefined) {
      frameIndex.value = frame;
    }
  }
}

const toBegin = () => {
  frameIndex.value = 1
}

const toEnd = () => {
  frameIndex.value = totalFrame.value
}

const togglePlay = () => {
  playing.value = !playing.value
  if (playing.value) {
    if (frameIndex.value === totalFrame.value && frameIndex.value !== 0) {
      frameIndex.value = 1;
    }
    const id = setInterval(() => {
      if (frameIndex.value < totalFrame.value) {
        frameIndex.value += 1
      } else {
        stopPlay()
      }
    }, 1000 / 30)
    playInterval.value = id;
  } else {
    stopPlay()
  }
}

const realVideoIndex = computed(() => {
  if (onlyShowUnmark.value) {
    return videoIndexes.value.filter(e => !markings.value.has(e))
  } else {
    return videoIndexes.value
  }
})

const keyFrameMarked = computed(() => {
  if (currVideoId.value != null) {
    return markings.value.has(currVideoId.value)
  }
})

const unmarked = computed(() => {
  const unmarked = videoIndexes.value.filter(e => !markings.value.has(e))
  return unmarked.sort()
})

onMounted(async () => {
  const videoIds: Array<number> = await invoke('get_video_id')
  videoIndexes.value = videoIds;

  const raw_label = await invoke('load_label') as string;
  const obj = JSON.parse(raw_label)
  const ret = new Map()
  for (const [key, value] of Object.entries(obj)) {
    ret.set(parseInt(key), value)
  }
  markings.value = ret;

})

const probeRealQuat = ref(Quaternion.Identity())

watch(frameIndex, async (curr, _) => {
  if (videoFrames.value !== null && curr <= videoFrames.value.length) {
    if (stage.value !== null) {
      const ctx = stage.value.getContext('2d')
      ctx?.drawImage(videoFrames.value[curr - 1], 0, 0)
    }
  }
  updateRealPos(curr)
  return
})

const setMark = () => {
  const videoId = currVideoId.value;
  if (videoId === null) {
    return;
  }
  const frameId = frameIndex.value;
  markings.value.set(videoId, frameId);
}

const saveMark = async () => {
  const marking = markings.value;
  const to_save: {[index: string]: number} = {};

  for (const [key, value] of marking.entries()) {
    to_save[key.toString()] = value
  }
  const encoded = JSON.stringify(to_save)
  await invoke('save_label', {"content": encoded})
}

const toNextVideo = () => {
  if (unmarked.value.length > 0) {
    updateVideoId(unmarked.value[0].toString())
  }
}

</script>

<template>
  <div class="container">

    <div class="row">
      <button @click="setMark" style="background-color: lightblue;">
        标记
      </button>

      <button @click="togglePlay" :disabled="totalFrame == 0">
        {{ playing ? "暂停" : "播放" }}
      </button>

      <button @click="toNextVideo" :class="{'enabled': unmarked.length > 0}">
        下一个视频
      </button>

      <button :class="{'enabled': keyFrameMarked}" @click="toKeyFrame">
        回关键帧
      </button>

      <button @click="toBegin">
        开头
      </button>
      <button @click="prevFrame" :disabled="frameIndex <= 1">
        上一帧
      </button>
      <button @click="nextFrame" :disabled="frameIndex == totalFrame">
        下一帧
      </button>
      <button @click="toEnd">
        结尾
      </button>

      <button @click="saveMark" style="background-color: lightgreen;">
        保存
      </button>

      <span v-if="keyFrameMarked">
        当前关键帧: {{  markings.get(currVideoId!)  }}
      </span>

      <div style="flex-grow: 1;"></div>

      <span>
        视频编号: {{ currVideoId }}
      </span>

      <span>
        剩余: {{ unmarked.length }}
      </span>

      <label for="only-unmark">
        仅显示未标注:
      </label>
      <input type="checkbox" id="only-unmark" v-model="onlyShowUnmark"/>

      <label for="video-id-input">
        视频序号:
      </label>
      <input id="video-id-input" list="video-id-list" @input="ev => updateVideoId((ev.target as HTMLInputElement).value)" />
      <datalist id="video-id-list">
        <option v-for="i of realVideoIndex.sort()" :value="i"></option>
      </datalist>
    </div>

    <div class="row">
      <p>帧数 {{ frameIndex }} / {{ totalFrame }} </p>
      <input class="slider" style="flex-grow: 1;" type="range" min="1" :max="totalFrame" :value="frameIndex">
    </div>

    <div id="bottom-container">
      <div>
        <canvas ref="stage" width="800" height="600"></canvas>
      </div>
      <ReviewScene :real-rotation="probeRealQuat" :target-rotation="targetPos" :pred-rotation="predRot" />
    </div>

  </div>
</template>

<style>
#bottom-container {
  width: 100%;
  display: grid;
  grid-template-columns: repeat(2, 1fr);
  flex-grow: 1;
  column-gap: 10px;
}

.big {
  width: 100%
}

.slider {
  outline: none;
}
</style>
